# -*- coding: utf-8 -*-

"""
Alphabet maps objects to integer ids. It provides two way mapping from the index to the objects.
"""
import json
import os


class Alphabet:
    def __init__(self, name, label=False, keep_growing=True):
        """
        初始化字母表类，创建一个对象与整数ID的映射表。
        :param name: 字母表的名称
        :param label: 是否是标签类型的字母表，默认为 False（如果是 False，表示字母表包含未知标记 "</unk>"）
        :param keep_growing: 是否允许字母表继续扩展新实例，默认为 True
        """
        self.__name = name
        self.UNKNOWN = "</unk>" # 设置默认未登录词的特殊标记为 "</unk>"
        self.label = label
        self.instance2index = {} # 初始化实例到索引的映射字典
        self.instances = []  # 实例列表
        self.keep_growing = keep_growing # 是否允许字母表扩展新实例

        # Index 0 is occupied by default, all else following.
        self.default_index = 0  # 默认索引为 0，表示特殊的 “UNKNOWN” 实例
        self.next_index = 1    # 从索引 1 开始，保持字母表索引递增
        if not self.label: # 如果不是标签字母表，默认添加一个 "UNKNOWN" 实例
            self.add(self.UNKNOWN)

    def clear(self, keep_growing=True):
        """
        清空字母表内容，并重置字母表的状态。
        :param keep_growing: 是否继续允许字母表扩展新实例，默认为 True
        """
        self.instance2index = {} # 清空映射字典
        self.instances = [] # 清空实例列表
        self.keep_growing = keep_growing # 设置是否允许继续扩展

        # Index 0 is occupied by default, all else following.
        # 重置索引信息
        self.default_index = 0
        self.next_index = 1

    def add(self, instance):
        """
        向字母表中添加一个新的实例（对象）。
        :param instance: 要添加的实例（对象）
        """
        if instance not in self.instance2index:
            self.instances.append(instance) # 将实例加入实例列表
            self.instance2index[instance] = self.next_index # 将实例与索引映射
            self.next_index += 1 # 更新下一个可用索引

    def get_index(self, instance):
        """
        获取一个实例的索引。
        :param instance: 要查找的实例
        :return: 实例对应的索引，如果实例不存在且允许扩展，则会添加实例并返回新索引
        """
        try:
            # 如果实例已经存在，返回对应的索引
            return self.instance2index[instance]
        except KeyError:
            # 如果实例不存在，并且允许扩展，则添加实例
            if self.keep_growing:
                index = self.next_index
                self.add(instance) # 添加新实例
                return index # 返回新实例的索引
            else:
                # 如果不允许扩展，返回 "UNKNOWN" 实例的索引
                return self.instance2index[self.UNKNOWN]

    def get_instance(self, index):
        """
        根据索引获取对应的实例（对象）。
        :param index: 实例对应的索引
        :return: 对应索引的实例，如果索引为 0 则返回 None，如果索引越界则返回第一个实例
        """
        if index == 0:
            # First index is occupied by the wildcard element.
            # 如果索引是 0，返回 None（这是为特殊的 "UNKNOWN" 保留的索引）
            return None
        try:
            # 尝试根据索引获取实例
            return self.instances[index - 1]
        except IndexError:
            # 如果索引超出范围，返回第一个实例，并打印警告
            print(
                "WARNING:Alphabet get_instance ,unknown instance index {}, return the first label.".format(
                    index
                )
            )
            return self.instances[0]

    def size(self):
        """
        返回字母表中实例的数量（包括 "UNKNOWN"）。
        :return: 字母表实例的数量
        """
        # if self.label:
        #     return len(self.instances)
        # else:
        return len(self.instances) + 1

    def iteritems(self):
        """
        返回实例到索引的所有映射项。
        :return: 实例到索引的映射（字典的键值对）
        """
        return self.instance2index.items()

    def enumerate_items(self, start=1):
        """
        从指定的索引位置开始返回实例的枚举。
        :param start: 开始的索引位置，默认为 1
        :return: 返回指定范围内的实例及其对应的索引
        """
        if start < 1 or start >= self.size():
            raise IndexError("Enumerate is allowed between [1 : size of the alphabet)")
        return zip(range(start, len(self.instances) + 1), self.instances[start - 1 :])

    def close(self):
        """
        禁止字母表继续扩展新的实例。
        """
        self.keep_growing = False

    def open(self):
        """
        允许字母表继续扩展新的实例。
        """
        self.keep_growing = True

    def get_content(self):
        """
        返回字母表的内容，包括实例到索引的映射和实例列表。
        :return: 包含字母表内容的字典
        """
        return {"instance2index": self.instance2index, "instances": self.instances}

    def from_json(self, data):
        """
        从 JSON 格式的数据加载字母表的内容。
        :param data: 包含字母表内容的 JSON 数据
        """
        self.instances = data["instances"]
        self.instance2index = data["instance2index"]

    def save(self, output_directory, name=None):
        """
        Save both alhpabet records to the given directory.
        :param output_directory: Directory to save model and weights.
        :param name: The alphabet saving name, optional.
        :return:
        将字母表保存到指定的目录中的 JSON 文件。
        :param output_directory: 保存字母表的目录路径
        :param name: 字母表保存时的文件名（如果未指定则使用字母表的名称）
        """
        saving_name = name if name else self.__name
        try:
            json.dump(
                self.get_content(),
                open(os.path.join(output_directory, saving_name + ".json"), "w"),
            )
        except Exception as e:
            print("Exception: Alphabet is not saved: " % repr(e))

    def load(self, input_directory, name=None):
        """
        Load model architecture and weights from the give directory. This allow we use old models even the structure
        changes.
        :param input_directory: Directory to save model and weights
        :return:
        从指定的目录加载保存的字母表 JSON 文件。
        :param input_directory: 存储字母表文件的目录路径
        :param name: 字母表文件名（如果未指定则使用字母表的名称）

        """
        loading_name = name if name else self.__name # 如果未指定文件名，使用字母表名称
        self.from_json(
            json.load(open(os.path.join(input_directory, loading_name + ".json")))
        )
